from flask import Flask, render_template, request, redirect, url_for, flash, session, jsonify, abort, \
    send_from_directory
from flask_wtf.csrf import CSRFProtect
from werkzeug.utils import secure_filename
import os
import sqlite3
from functools import wraps
from datetime import timedelta

# 初始化应用
app = Flask(__name__)
app.secret_key = 'your-very-secret-key-here'  # 请替换为随机密钥
app.config['UPLOAD_FOLDER'] = os.path.abspath('upload')
app.config['PERMANENT_SESSION_LIFETIME'] = timedelta(hours=1)
os.makedirs(app.config['UPLOAD_FOLDER'], exist_ok=True)

# 启用CSRF保护
csrf = CSRFProtect(app)


# 数据库连接辅助函数
def get_db():
    conn = sqlite3.connect('student_fileserver.db')
    conn.row_factory = sqlite3.Row
    return conn


# 登录检查装饰器
def login_required(f):
    @wraps(f)
    def decorated_function(*args, **kwargs):
        if 'username' not in session:
            if request.headers.get('X-Requested-With') == 'XMLHttpRequest':
                return jsonify({'error': '请先登录'}), 401
            flash('请先登录', 'warning')
            return redirect(url_for('login', next=request.url))
        return f(*args, **kwargs)

    return decorated_function


# 获取用户上传目录
def get_user_upload_dir(username):
    user_dir = os.path.join(app.config['UPLOAD_FOLDER'], secure_filename(username))
    os.makedirs(user_dir, exist_ok=True)
    return user_dir


# 路由定义
@app.route('/')
def index():
    if 'username' in session:
        return redirect(url_for('view'))
    return redirect(url_for('login'))


@app.route('/login/', methods=['GET', 'POST'])
def login():
    if request.method == 'POST':
        username = request.form.get('username', '').strip()
        password = request.form.get('password', '')

        if not username or not password:
            if request.headers.get('X-Requested-With') == 'XMLHttpRequest':
                return jsonify({'error': '用户名和密码不能为空'}), 400
            flash('用户名和密码不能为空', 'error')
            return redirect(url_for('login'))

        with get_db() as conn:
            user = conn.execute(
                "SELECT id, username, password FROM student_fileserver_usr WHERE username = ?",
                (username,)
            ).fetchone()

        if user and user['password'] == password:
            session['username'] = username
            session['user_id'] = user['id']

            if request.headers.get('X-Requested-With') == 'XMLHttpRequest':
                return jsonify({'success': '登录成功', 'redirect': url_for('view')})

            flash('登录成功', 'success')
            return redirect(url_for('view'))
        else:
            if request.headers.get('X-Requested-With') == 'XMLHttpRequest':
                return jsonify({'error': '用户名或密码错误'}), 401
            flash('用户名或密码错误', 'error')

    return render_template('login.html')


@app.route('/logout/')
def logout():
    session.clear()
    flash('您已安全退出', 'info')
    return redirect(url_for('login'))


@app.route('/view/')
@login_required
def view():
    with get_db() as conn:
        files = conn.execute(
            "SELECT id, addr FROM student_fileserver WHERE username = ?",
            (session['username'],)
        ).fetchall()

    if request.headers.get('X-Requested-With') == 'XMLHttpRequest':
        return jsonify({
            'files': [dict(file) for file in files],
            'username': session['username']
        })
    return render_template('view.html', files=files)


@app.route('/upload/<filename>')
@login_required
def download_file(filename):
    with get_db() as conn:
        file_exists = conn.execute(
            "SELECT 1 FROM student_fileserver WHERE addr = ? AND username = ?",
            (filename, session['username'])
        ).fetchone()

    if not file_exists:
        abort(404)

    return send_from_directory(
        directory=get_user_upload_dir(session['username']),
        path=secure_filename(filename),
        as_attachment=True
    )


@app.route('/add/', methods=['GET', 'POST'])
@login_required
def upload_file():
    if request.method == 'POST':
        if 'file' not in request.files:
            flash('未选择文件', 'error')
            return redirect(request.url)

        file = request.files['file']
        if file.filename == '':
            flash('未选择文件', 'error')
            return redirect(request.url)

        try:
            filename = secure_filename(file.filename)
            if not filename:
                raise ValueError("无效文件名")

            user_dir = get_user_upload_dir(session['username'])
            dest_path = os.path.join(user_dir, filename)

            if os.path.exists(dest_path):
                flash('文件已存在', 'warning')
                return redirect(request.url)

            file.save(dest_path)

            with get_db() as conn:
                conn.execute(
                    "INSERT INTO student_fileserver (addr, username) VALUES (?, ?)",
                    (filename, session['username'])
                )
                conn.commit()

            flash('上传成功', 'success')
            return redirect(url_for('view'))

        except Exception as e:
            flash(f'上传失败: {str(e)}', 'error')

    return render_template('add.html')


@app.route('/delete/<int:file_id>', methods=['POST'])
@login_required
def delete_file(file_id):
    try:
        with get_db() as conn:
            # 获取文件名
            file_record = conn.execute(
                "SELECT addr FROM student_fileserver WHERE id = ? AND username = ?",
                (file_id, session['username'])
            ).fetchone()

            if not file_record:
                flash('文件不存在或无权操作', 'error')
                return redirect(url_for('view'))

            filename = file_record['addr']
            filepath = os.path.join(
                get_user_upload_dir(session['username']),
                secure_filename(filename)
            )

            # 删除数据库记录
            conn.execute(
                "DELETE FROM student_fileserver WHERE id = ? AND username = ?",
                (file_id, session['username'])
            )

            # 删除物理文件
            if os.path.exists(filepath):
                os.remove(filepath)

            conn.commit()
            flash('删除成功', 'success')

    except Exception as e:
        flash(f'删除失败: {str(e)}', 'error')

    return redirect(url_for('view'))


# 初始化数据库
def init_db():
    with sqlite3.connect("student_fileserver.db") as conn:
        cursor = conn.cursor()
        cursor.execute("""
            CREATE TABLE IF NOT EXISTS student_fileserver (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                addr TEXT NOT NULL,
                username TEXT NOT NULL,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                UNIQUE(addr, username)
            )
        """)
        cursor.execute("""
            CREATE TABLE IF NOT EXISTS student_fileserver_usr (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                username TEXT UNIQUE NOT NULL,
                password TEXT NOT NULL
            )
        """)
        conn.commit()


if __name__ == '__main__':
    init_db()
    app.run(host='0.0.0.0', port=5000, debug=True)